package com.zzyang.mybatis.jdbc.core;

import com.zzyang.mybatis.jdbc.config.IConstants;
import com.zzyang.mybatis.jdbc.enums.DataSourceType;
import com.zzyang.mybatis.jdbc.loadbalance.LoadBalance;
import com.zzyang.mybatis.jdbc.loadbalance.RandomLoadBalance;
import com.zzyang.mybatis.jdbc.loadbalance.RoundRobinLoadBalance;
import com.zzyang.mybatis.jdbc.spring.SpringContextUtils;
import com.zzyang.mybatis.jdbc.spring.config.SeparationJdbcProperties;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.Map;

/**
 * 默认实现数据源路由方式
 *
 * @author zhuoziyang
 * @date 2021/4/10 09:04
 */
public class DefaultRoutingDataSourceKeyDecider implements RoutingDataSourceKeyDecider, InitializingBean {

    @Autowired
    private DataSourceContext dataSourceContext;

    private LoadBalance loadBalance;

    private boolean onlyUseMaster = false;

    private String masterDataSourceId;

    @Override
    public String decideRoutingKey() {
        if (onlyUseMaster) {
            return masterDataSourceId;
        }

        DataSourceType dataSourceType = DataSourceSynchronizationManager.getDataSourceType();
        if (dataSourceType == null || DataSourceType.MASTER.equals(dataSourceType)) {
            return masterDataSourceId;
        }

        String slaveDataSourceId = loadBalance.select(dataSourceContext.getAvailableSlaveDataSourceIds());
        if (IConstants.NO_AVAILABLE.equals(slaveDataSourceId)) {
            return masterDataSourceId;
        }
        DataSourceSynchronizationManager.setSlaveDataSourceId(slaveDataSourceId);
        return slaveDataSourceId;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        masterDataSourceId = dataSourceContext.getMasterDataSourceId();

        Map<String, DataSource> slaveDataSource = dataSourceContext.getSlaveDataSource();
        if (CollectionUtils.isEmpty(slaveDataSource)) {
            onlyUseMaster = true;
        }

        SeparationJdbcProperties properties = SpringContextUtils.getBean(SeparationJdbcProperties.class);
        initLoadBalance(properties.getLoadbalance());
    }

    private void initLoadBalance(String loadbalance) {
        if (onlyUseMaster) {
            return;
        }
        if (!StringUtils.hasText(loadbalance) || IConstants.RANDOM_LOAD_BALANCE.equals(loadbalance)) {
            loadBalance = new RandomLoadBalance();
        } else if (IConstants.ROUND_ROBIN_LOAD_BALANCE.equals(loadbalance)) {
            loadBalance = new RoundRobinLoadBalance();
        } else {
            try {
                loadBalance = SpringContextUtils.getBean(loadbalance, LoadBalance.class);
            } catch (Exception e) {
                throw new IllegalArgumentException("无法根据 loadbalance=" + loadbalance + " 获取负载均衡策略");
            }
        }
    }

}
